import asyncio
import json
import logging
from typing import Dict, Optional

import websockets
from pydantic import BaseModel

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


# 定义MCP消息模型
class MCPMessage(BaseModel):
    type: str
    content: Dict
    id: Optional[str] = None


class MCPTool(BaseModel):
    name: str
    description: str
    parameters: Dict


# Ollama API配置
OLLAMA_API_URL = "http://localhost:11434/api"
MODEL_NAME = "qwen:8b"

# 示例工具定义
TOOLS = [
    MCPTool(
        name="calculator",
        description="A simple calculator that can perform basic arithmetic operations",
        parameters={
            "type": "object",
            "properties": {
                "operation": {
                    "type": "string",
                    "enum": ["add", "subtract", "multiply", "divide"]
                },
                "a": {"type": "number"},
                "b": {"type": "number"}
            },
            "required": ["operation", "a", "b"]
        }
    ),
    MCPTool(
        name="weather",
        description="Get current weather information for a location",
        parameters={
            "type": "object",
            "properties": {
                "location": {"type": "string"},
                "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]}
            },
            "required": ["location"]
        }
    )
]


class MCPClient:
    def __init__(self, server_url: str = "ws://localhost:8000/ws"):
        self.server_url = server_url
        self.ws = None
        self.tools = []
        self.connected = False
        self.reconnect_delay = 1  # 初始重连延迟（秒）
        self.max_reconnect_delay = 30  # 最大重连延迟（秒）
        self.max_retries = 3  # 最大重试次数
        self.retry_count = 0

    async def connect(self):
        """连接到MCP服务器"""
        while self.retry_count < self.max_retries:
            try:
                logger.info(
                    f"Connecting to MCP server at {self.server_url} (attempt {self.retry_count + 1}/{self.max_retries})")
                self.ws = await websockets.connect(
                    self.server_url,
                    ping_interval=20,  # 每20秒发送一次ping
                    ping_timeout=10,  # ping超时时间为10秒
                    close_timeout=5  # 关闭超时时间为5秒
                )
                self.connected = True
                self.reconnect_delay = 1  # 重置重连延迟
                self.retry_count = 0  # 重置重试计数
                logger.info("Connected to MCP server")

                # 接收工具定义
                try:
                    response = await asyncio.wait_for(self.ws.recv(), timeout=10)
                    tools_data = json.loads(response)

                    if tools_data["type"] == "tools":
                        self.tools = tools_data["content"]["tools"]
                        print("\nAvailable tools:")
                        for tool in self.tools:
                            print(f"- {tool['name']}: {tool['description']}")
                        return True
                    elif tools_data["type"] == "error":
                        logger.error(f"Server error: {tools_data['content']['message']}")
                        raise ConnectionError(tools_data['content']['message'])
                except asyncio.TimeoutError:
                    logger.error("Timeout waiting for tools definition")
                    raise ConnectionError("Timeout waiting for tools definition")

            except websockets.exceptions.ConnectionClosed as e:
                logger.warning(f"Connection closed: {e}")
                self.connected = False
                self.retry_count += 1
            except Exception as e:
                logger.error(f"Connection error: {e}")
                self.connected = False
                self.retry_count += 1

            if self.retry_count < self.max_retries:
                # 等待重连
                wait_time = min(self.reconnect_delay * (2 ** self.retry_count), self.max_reconnect_delay)
                logger.info(f"Waiting {wait_time} seconds before retrying...")
                await asyncio.sleep(wait_time)
            else:
                logger.error("Max retry attempts reached")
                raise ConnectionError("Failed to connect after maximum retry attempts")

    async def ensure_connection(self):
        """确保连接状态"""
        if not self.connected or not self.ws:
            await self.connect()

    async def send_message(self, message: str):
        """发送消息到服务器"""
        await self.ensure_connection()

        try:
            await asyncio.wait_for(
                self.ws.send(json.dumps({
                    "type": "chat",
                    "content": {"message": message}
                })),
                timeout=5
            )
            logger.debug(f"Sent message: {message}")
        except asyncio.TimeoutError:
            logger.error("Timeout sending message")
            self.connected = False
            raise
        except Exception as e:
            logger.error(f"Error sending message: {e}")
            self.connected = False
            raise

    async def call_tool(self, tool_name: str, parameters: Dict):
        """调用工具"""
        await self.ensure_connection()

        try:
            await asyncio.wait_for(
                self.ws.send(json.dumps({
                    "type": "tool_call",
                    "content": {
                        "name": tool_name,
                        "parameters": parameters
                    }
                })),
                timeout=5
            )
            logger.debug(f"Sent tool call: {tool_name} with parameters: {parameters}")
        except asyncio.TimeoutError:
            logger.error("Timeout calling tool")
            self.connected = False
            raise
        except Exception as e:
            logger.error(f"Error calling tool: {e}")
            self.connected = False
            raise

    async def receive_response(self):
        """接收服务器响应"""
        await self.ensure_connection()

        try:
            response = await asyncio.wait_for(self.ws.recv(), timeout=30)
            response_data = json.loads(response)
            logger.debug(f"Received response: {response_data}")

            if response_data["type"] == "error":
                logger.error(f"Server error: {response_data['content']['message']}")
                raise Exception(response_data['content']['message'])

            return response_data
        except asyncio.TimeoutError:
            logger.error("Timeout receiving response")
            self.connected = False
            raise
        except websockets.exceptions.ConnectionClosed:
            logger.warning("Connection closed while receiving response")
            self.connected = False
            raise
        except Exception as e:
            logger.error(f"Error receiving response: {e}")
            self.connected = False
            raise

    async def close(self):
        """关闭连接"""
        if self.ws:
            try:
                await asyncio.wait_for(self.ws.close(), timeout=5)
                logger.info("Connection closed")
            except Exception as e:
                logger.error(f"Error closing connection: {e}")
            finally:
                self.ws = None
                self.connected = False


async def main():
    client = MCPClient()
    try:
        await client.connect()

        while True:
            try:
                # 获取用户输入
                user_input = input("\nEnter your message (or 'quit' to exit): ").strip()
                if user_input.lower() == 'quit':
                    break

                # 发送消息
                await client.send_message(user_input)

                # 接收响应
                while True:
                    try:
                        response = await client.receive_response()
                        if response["type"] == "response":
                            print(f"\nBot: {response['content']['message']}")
                            break
                        elif response["type"] == "tool_response":
                            print(f"\nTool Response: {json.dumps(response['content'], indent=2)}")
                            # 如果是工具响应，继续等待模型响应
                            continue
                    except Exception as e:
                        logger.error(f"Error processing response: {e}")
                        if not client.connected:
                            print("\nConnection lost. Attempting to reconnect...")
                            try:
                                await client.connect()
                                # 重新发送消息
                                await client.send_message(user_input)
                                continue
                            except Exception as reconnect_error:
                                logger.error(f"Reconnection failed: {reconnect_error}")
                                break
                        break

            except Exception as e:
                logger.error(f"Error in main loop: {e}")
                if not client.connected:
                    print("\nConnection lost. Attempting to reconnect...")
                    try:
                        await client.connect()
                    except Exception as reconnect_error:
                        logger.error(f"Reconnection failed: {reconnect_error}")
                        break

    except KeyboardInterrupt:
        print("\nExiting...")
    except Exception as e:
        logger.error(f"Fatal error: {e}")
    finally:
        await client.close()


if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        print("\nProgram terminated by user")
    except Exception as e:
        logger.error(f"Program error: {e}")
